import numpy as np
import torch 
import torch.nn as nn 
from monai.networks.nets.basic_unet import BasicUNet
from monai.networks.nets.unetr import UNETR
from monai.networks.nets.swin_unetr import SwinUNETR
from monai.inferers import SlidingWindowInferer
from light_training.evaluation.metric import dice
from light_training.trainer import Trainer
from monai.utils import set_determinism
from light_training.utils.lr_scheduler import LinearWarmupCosineAnnealingLR
from light_training.utils.files_helper import save_new_model_and_delete_last
from models.unet_soft_loss.uent2d import UNet2D
from models.unet_soft_loss.uent3d import UNet3D
from models.unet_soft_loss.uent25d import UNet25D
from models.att_unet_ml.att_unet25d import AttUNet25D
from models.att_unet_ml.att_unet2d import AttUNet2D
from models.att_unet_ml.att_unet3d import AttUNet3D
from monai.networks.nets.segresnet import SegResNet
from models.transbts.TransBTS_downsample8x_skipconnection import TransBTS
from models.nestedformer.nested_former import NestedFormer
from models.swinunet2d.swinunet import SwinUnet
from einops import rearrange
from monai.networks.nets.vnet import VNet
from models.modelgenesis.unet3d import UNet3DModelGen
from models.transvw.models.ynet3d import UNet3DTransVW
from monai.networks.nets.attentionunet import AttentionUnet
import random
import h5py
import SimpleITK as sitk 
from torch.utils.data import Dataset
from models.hdensenet_huanhu.net import dense_rnn_net
from models.unet_nested.unet_nested_3d import UNet_Nested3D
import glob 
from utils.utils import SoftmaxLoss
from monai.utils import set_determinism
set_determinism(1)

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

# logdir = "./logs_huanhu/unet_ml_e1000_schdule"
# logdir = "./logs/unet_no_ml"
logdir = "./logs/att_unet_no_ml"

env = "pytorch"
model_save_path = os.path.join(logdir, "model")
max_epoch = 100
batch_size = 1
val_every = 5
num_gpus = 1
device = "cuda:1"

uncer_threshold = 0.5

images_paths = sorted(glob.glob("/home/xingzhaohu/sharefs/datasets/huanhu/IBSR_3/ii_IBSR_*_ana_strip.nii.gz"))
seg_paths = sorted(glob.glob("/home/xingzhaohu/sharefs/datasets/huanhu/IBSR_3/IBSR_*_segTRI_ana.nii.gz"))
print(images_paths)
print(seg_paths)

def compute_uncer(pred_out):

    uncer_out = torch.softmax(pred_out, dim=1)
    ## 计算学习比重
    uncer_out = torch.sum(-uncer_out * torch.log(uncer_out), dim=1, keepdim=True)

    return uncer_out

class Dataset3d(Dataset):
    """
    """
    def __init__(
        self,
        image_paths,
        label_paths,
        has_tqdm=True
    ) -> None:
        """
        Args:
            data: input data to load and transform to generate dataset for model.
            transform: transforms to execute operations on input data.
            cache_num: number of items to be cached. Default is `sys.maxsize`.
                will take the minimum of (cache_num, data_length x cache_rate, data_length).
            cache_rate: percentage of cached data in total, default is 1.0 (cache all).
                will take the minimum of (cache_num, data_length x cache_rate, data_length).
            num_workers: the number of worker threads to use.
                If 0 a single thread will be used. Default is 0.
        """

        super(Dataset3d, self).__init__()
        self.image_paths = image_paths
        self.label_paths = label_paths

    def _load_cache_item(self, image_path, label_path):
        img = sitk.ReadImage(image_path)
        img = sitk.GetArrayFromImage(img)
        # img = img.transpose(1, 0, 2)
        label = sitk.ReadImage(label_path)
        label = sitk.GetArrayFromImage(label)
        start = 0
        end = -1
        for i in range(label.shape[0]):
            if label[i].sum() != 0:
                start = i
                break
        for j in range(label.shape[0]-1, 0, -1):
            if label[j].sum() != 0:
                end = j
                break
        img = img[start:end+1]
        label = label[start:end+1]

        # print("start is {}, end is {}".format(start, end))
        # label = label.transpose(1, 0, 2)
        return img, label

    def __getitem__(self, index):

        image, label = self._load_cache_item(self.image_paths[index], self.label_paths[index])

        image = np.expand_dims(image, axis=0)
        image = (image - np.mean(image)) / np.std(image)

        return image.astype(np.float32), label.astype(np.uint8)

    def __len__(self):
        return len(self.image_paths)

class BraTSTrainer(Trainer):
    def __init__(self, env_type, max_epochs, batch_size, device="cpu", val_every=1, num_gpus=1, logdir="./logs/", master_ip='localhost', master_port=17750, training_script="train.py",fold=0):
        super().__init__(env_type, max_epochs, batch_size, device, val_every, num_gpus, logdir, master_ip, master_port, training_script)
    
        self.fold=fold
        self.best_mean_dice3d = 0.0
        self.best_mean_dice25d = 0.0
        self.best_mean_dice2d = 0.0
        self.auto_optim = False

        self.model3d = AttUNet3D(in_ch=1, out_ch=4)
        self.model25d = AttUNet25D(in_ch=3, out_ch=4)
        self.model2d = AttUNet2D(in_ch=1, out_ch=4)
        self.optimizer3d = torch.optim.Adam(self.model3d.parameters(), lr=1e-4, weight_decay=1e-5)
        self.optimizer25d = torch.optim.Adam(self.model25d.parameters(), lr=1e-4, weight_decay=1e-5)
        self.optimizer2d = torch.optim.Adam(self.model2d.parameters(), lr=1e-4, weight_decay=1e-5)
        
        # self.scheduler3d = LinearWarmupCosineAnnealingLR(self.optimizer3d,
        #                                           warmup_epochs=10,
        #                                           max_epochs=max_epochs,
        #                                           warmup_start_lr=1e-6)
        
        # self.scheduler25d = LinearWarmupCosineAnnealingLR(self.optimizer25d,
        #                                           warmup_epochs=10,
        #                                           max_epochs=max_epochs,
        #                                           warmup_start_lr=1e-6)
        
        # self.scheduler2d = LinearWarmupCosineAnnealingLR(self.optimizer2d,
        #                                           warmup_epochs=10,
        #                                           max_epochs=max_epochs,
        #                                           warmup_start_lr=1e-6)
        
        self.loss_func = SoftmaxLoss()
        self.loss_mse = nn.MSELoss()

    def training_step(self, batch):
        import time
        self.model3d.train()
        self.model25d.train()
        self.model2d.train()
        image, label = self.get_input(batch)
        self.model3d.to(self.device)
        self.model25d.to(self.device)
        self.model2d.to(self.device)

        pred3d = self.model3d(image)
        pred25d = self.model25d(image)
        pred2d = self.model2d(image)

        hard_loss_3d = self.loss_func(pred3d, label)
        hard_loss_25d = self.loss_func(pred25d, label)
        hard_loss_2d = self.loss_func(pred2d, label)

        loss3d = hard_loss_3d

        loss25d = hard_loss_25d

        loss2d = hard_loss_2d

        loss3d.backward()
        loss25d.backward()
        loss2d.backward()

        self.optimizer3d.step()
        self.optimizer25d.step()
        self.optimizer2d.step()

        self.optimizer3d.zero_grad()
        self.optimizer25d.zero_grad()
        self.optimizer2d.zero_grad()
    
        if self.global_step % 100 == 0:
            print(loss3d, loss25d, loss2d)


    def get_input(self, batch):
        image, label = batch 
        
        image = nn.functional.interpolate(image, size=(64, 128, 256), mode="trilinear", align_corners=False)
        label = torch.unsqueeze(label, dim=1)
        label = nn.functional.interpolate(label, size=(64, 128, 256), mode="nearest")
        label = torch.squeeze(label, dim=1).long()

        return image, label 

    def validation_step(self, batch):
        self.model3d.eval()
        self.model25d.eval()
        self.model2d.eval()

        image, label = self.get_input(batch)
        self.model3d.to(self.device)
        self.model25d.to(self.device)
        self.model2d.to(self.device)

        output3d = self.model3d(image)
        output25d = self.model25d(image)
        output2d = self.model2d(image)

        output3d = output3d.argmax(dim=1).cpu().numpy()
        output25d = output25d.argmax(dim=1).cpu().numpy()
        output2d = output2d.argmax(dim=1).cpu().numpy()

        target = label.cpu().numpy()

        o = output3d == 1; t = target == 1  # ce
        csf = dice(o, t)
        o = output3d == 2; t = target == 2  # ce
        gm = dice(o, t)
        o = output3d == 3; t = target == 3  # ce
        wm = dice(o, t)
        dice_score3d = (csf + gm + wm) / 3

        o = output25d == 1; t = target == 1  # ce
        csf = dice(o, t)
        o = output25d == 2; t = target == 2  # ce
        gm = dice(o, t)
        o = output25d == 3; t = target == 3  # ce
        wm = dice(o, t)
        dice_score25d = (csf + gm + wm) / 3

        o = output2d == 1; t = target == 1  # ce
        csf = dice(o, t)
        o = output2d == 2; t = target == 2  # ce
        gm = dice(o, t)
        o = output2d == 3; t = target == 3  # ce
        wm = dice(o, t)
        dice_score2d = (csf + gm + wm) / 3
        
        return [dice_score3d, dice_score25d, dice_score2d]

    def validation_end(self, mean_val_outputs):
        dice_score3d, dice_score25d, dice_score2d = mean_val_outputs

        self.log("dice_score_3d", dice_score3d, step=self.epoch)
        self.log("dice_score_25d", dice_score25d, step=self.epoch)
        self.log("dice_score_2d", dice_score2d, step=self.epoch)
        if self.best_mean_dice3d < dice_score3d:
            self.best_mean_dice3d = dice_score3d
            save_new_model_and_delete_last(self.model3d, 
                                        os.path.join(model_save_path, 
                                        f"best_model3d_{dice_score3d:.4f}_{self.fold}.pt"), 
                                        delete_symbol=f"best_model3d_*_{self.fold}")
        save_new_model_and_delete_last(self.model3d, 
                                        os.path.join(model_save_path, 
                                        f"final_model3d_{dice_score3d:.4f}_{self.fold}.pt"), 
                                        delete_symbol=f"final_model3d_*_{self.fold}")
        
        if self.best_mean_dice25d < dice_score25d:
            self.best_mean_dice25d = dice_score25d

            save_new_model_and_delete_last(self.model25d, 
                                        os.path.join(model_save_path, 
                                        f"best_model25d_{dice_score25d:.4f}_{self.fold}.pt"), 
                                        delete_symbol=f"best_model25d_*_{self.fold}")

        save_new_model_and_delete_last(self.model25d, 
                                        os.path.join(model_save_path, 
                                        f"final_model25d_{dice_score25d:.4f}_{self.fold}.pt"), 
                                        delete_symbol=f"final_model25d_*_{self.fold}")

        if self.best_mean_dice2d < dice_score2d:
            self.best_mean_dice2d = dice_score2d

            save_new_model_and_delete_last(self.model2d, 
                                        os.path.join(model_save_path, 
                                        f"best_model2d_{dice_score2d:.4f}_{self.fold}.pt"), 
                                        delete_symbol=f"best_model2d_*_{self.fold}")
        save_new_model_and_delete_last(self.model2d, 
                                        os.path.join(model_save_path, 
                                        f"final_model2d_{dice_score2d:.4f}_{self.fold}.pt"), 
                                        delete_symbol=f"final_model2d_*_{self.fold}")
        
        print(f"dice_score3d is {dice_score3d}, dice_score25d is {dice_score25d}, dice_score2d is {dice_score2d}")

if __name__ == "__main__":

    from sklearn.model_selection import KFold  ## K折交叉验证

    X = np.arange(18)
    kfold = KFold(n_splits=6, shuffle=False)  ## kfold为KFolf类的一个对象
    fold = 0
    for a, b in kfold.split(X):  ## .split(X)方法返回迭代器，迭代器每次产生两个元素，1、训练数据集的索引；
        fold += 1
       
        print('Train_index: ', a, 'Validation_index:', b)
        # 构建训练集和验证集
        img_paths_train = []
        img_paths_val = []
        seg_paths_train = []
        seg_paths_val = []
        for i in a:
            img_paths_train.append(images_paths[i])
            seg_paths_train.append(seg_paths[i])
        for j in b:
            img_paths_val.append(images_paths[j])
            seg_paths_val.append(seg_paths[j])

        train_paths = {"image": img_paths_train, "label": seg_paths_train}
        val_paths = {"image": img_paths_val, "label": seg_paths_val}

        train_ds = Dataset3d(train_paths["image"], train_paths["label"])
        val_ds = Dataset3d(val_paths["image"], val_paths["label"])

        trainer = BraTSTrainer(env_type=env,
                                max_epochs=max_epoch,
                                batch_size=batch_size,
                                device=device,
                                logdir=logdir,
                                val_every=val_every,
                                num_gpus=num_gpus,
                                master_port=17751,
                                training_script=__file__,
                                fold=fold,
                                )

        trainer.train(train_dataset=train_ds, val_dataset=val_ds)

        print(f"第{fold}折训练完成。。。。")